import yaml


class DeepLearningYaml(yaml.YAMLObject):
    def __init__(self):
        super(DeepLearningYaml, self).__init__()
        self.batch_size = 8
        self.image_size = 512
        self.num_workers = 4
        self.lr = 0.001
        self.dataset = None
        self.gpus = [0]
        self.dir = None
        self.epochs = 100
        self.momentum = 0.9
        self.weight_decay = 0.0004
        self.wandb = False